feat(speechlm2): add padding-free packed speech training by pzelasko · Pull Request #16194 · NVIDIA-NeMo/Speech

pzelasko · 2026-09-01T20:56:44Z

feat(speechlm2): add padding-free packed speech training

Stack 3/7: 1/7 Encoder foundations → 2/7 Indexed/resumable data → Packed speech training → 4/7 Portable export + target-only vLLM → 5/7 MTP serving → 6/7 DFlash serving → 7/7 CCE dependency

Base: main

Head: codex/pr16158-rewrite/03-packed-speech-training

Tip: 919c31b5bf65f83e27e57cbe1f853cf06a8090e4

Summary

Add native THD Transformer/PEE execution and packed activation, waveform, feature, and SpecAugment paths.
Add padding-free SALM batches with CP/FSDP/FP8 alignment and activation checkpointing.
Add PEE execution controls, fused loss, and the FSDP-shared MTP training loss head.
Add optimizer-resume, garbage-collection, and OOMptimizer runtime hardening.
Keep the obsolete GGEMM/MoE implementations, benchmarks, imports, and tests removed throughout this layer.

Reviewer guide

The review question is: does SALM consume the encoder and data foundations correctly under packed, distributed, resumable training?

Index-format and routing implementations remain in PR 2. Checkpoint export and all vLLM behavior remain outside this PR.

Attribution

The replay preserves original authors, dates, source-commit references, and author-matching DCO sign-offs.

Validation

git diff --check: passed.
Compilation: passed.
Focused stack suite: 403 passed, 20 skipped.
Nithin review coverage: 207 packed ASR/SpeechLM2 tests passed; the packed Transformer GPU suite passed 53/53 on GB300.
Packed perception regression coverage passes on CPU (9 passed, 2 CUDA-only skipped) and under the GPU CI collection environment on GB300 (11 passed).
Removed the redundant Docker-only cut-cross-entropy metadata/import assertion; the optional code paths remain usable without the dependency, while its declaration and lock metadata are intentionally deferred to PR 7.
The FlashAttention device/provider cache regression test passes, including repeated dispatch on the same CUDA device.
All newly introduced packed-attention entry points and private helpers now document their contracts and behavior.
Rebuilt top-of-stack vLLM/export regression suite: 188 passed, 3 skipped.
Branch-scoped and complete-stack pre-commit: passed.
uv lock --check --offline: passed.
Range-diff confirms unchanged commits are patch-identical; the only deviations are the intentional obsolete-component cleanup and its downstream references.
Documentation is unchanged because this is a private helper extraction and runtime dispatch optimization with no public API or configuration change.

copy-pr-bot · 2026-09-01T20:56:48Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

pzelasko · 2026-09-01T22:49:24Z

/ok to test 42a1f6e

pzelasko · 2026-09-01T23:11:48Z

/ok to test a0eb08d

pzelasko · 2026-09-02T12:57:36Z

/ok to test 7eb5294

pzelasko · 2026-09-10T20:05:14Z

/ok to test 0a46d83

nithinraok · 2026-09-10T21:28:06Z

+    return flash_attn_varlen_func
+
+
+def _forward_sequence_packed_layer(layer, x, **kwargs):


are we adding these utils to fastconformer as well?

Another way of asking: Are we supporting packed sequence feature to FastConformer based encoders as well?

Not currently, but nothing stops us from backporting later. We should IMO.

nithinraok · 2026-09-10T21:29:17Z

        out = out.transpose(1, 2).contiguous().view(B, num_cur, self.d_model)
        return self.out_proj(out)

+    def forward_sequence_packed(


many doesn;t have doc strings, add wherever required.

Addressed in 25c5e7b. I documented all eight newly introduced packed private methods in transformer_encoder.py, expanded both public forward_sequence_packed() entry points with their argument and return contracts, and documented the extracted helper module plus every top-level helper and the nested ATen adapter. I kept the change documentation-only and limited to the newly introduced packed-attention code. The affected ASR API page now passes a warning-as-error Sphinx build, and the packed Transformer suite remains green.

pzelasko · 2026-09-11T01:08:15Z

/ok to test 25c5e7b

nithinraok

Thanks. LGTM

nithinraok · 2026-09-11T04:29:21Z

+python - <<'PY'
+from importlib.metadata import version
+
+import cut_cross_entropy


why is this required?

Big memory savings by avoiding logit materialization in CE loss.

Thanks, that explains why we use it. My question was about the version check in the Dockerfile: can we pin this in pyproject/uv.lock instead?

Agreed — removed in 3d8530f. Both compiled extras already pin cut-cross-entropy 25.3.2, pyproject.toml pins its exact source revision, and uv.lock plus uv sync --locked enforce that installation. The deleted import/version assertion duplicated those guarantees and did not exercise the fused GPU kernel.

can we do same for automodel version check.

pzelasko · 2026-09-11T14:58:00Z

/ok to test eb0d8f6

pzelasko · 2026-09-11T15:07:49Z

/ok to test 3d8530f

chtruong814

We'll need to review the inclusion of cut-cross-entropy.

chtruong814 · 2026-09-11T15:55:59Z

I recommend treating cut-cross-entropy as optional and not formally adding into pyproject.toml for now until review is done. So you can merge code and be unblocked.

pzelasko · 2026-09-11T16:31:11Z

/ok to test 7dc7739

github-actions · 2026-09-11T19:38:16Z

[🤖]: Hi @pzelasko 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

Introduce PackedEncoderActivations and native THD paths for waveform preprocessing, normalization, SpecAugment, Transformer/MoE/GGEMM/phPEE execution, activation checkpointing, and packed encoder benchmarks. Co-authored-by: Taejin Park <tango4j@gmail.com> Original-Commit: 9cd9a6a Original-Commit: 441eff6 Original-Commit: fcb1726 Original-Commit: 3ca6b8b Original-Commit: 56d9b52 Original-Commit: d49ae3c Original-Commit: 58b923d Original-Commit: d665c44 Original-Commit: 2fb5e1a Original-Commit: 009d1c1 Original-Commit: e5374ff Original-Commit: 4e034d0 Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Add native packed perception capabilities, independently chunked ASR and speaker encoders, packed FSDP gradient handling, PEE execution controls, and the speaker-encoder artifact renderer. Original-Commit: 6593f11 Original-Commit: 5ab986f Original-Commit: be3a0e9 Original-Commit: 4e034d0 Original-Commit: e5374ff Original-Commit: 8b8ea97 Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Carry packed waveform and sequence metadata through SALM batching, CP/FSDP/FP8 execution, checkpointing, fused CE and shared MTP loss heads; add text-only and missing-RTTM handling, lazy optimizer resume, step-aligned GC, timeout/OOM tuning, validator forwarding, configs, dependencies, documentation, and focused coverage. Co-authored-by: Kunal Dhawan <kunaldhawan97@gmail.com> Original-Commit: 29f8260 Original-Commit: 210a5b6 Original-Commit: c1a633a Original-Commit: 3355af1 Original-Commit: 303f36e Original-Commit: e1f2309 Original-Commit: ba97a73 Original-Commit: 95a77c3 Original-Commit: 6ba2158 Original-Commit: 2c645f4 Original-Commit: 556db89 Original-Commit: 5d68cbb Original-Commit: e376f26 Original-Commit: 3d7a2b8 Original-Commit: 27cac7e Original-Commit: b0fbee2 Original-Commit: 6689df1 Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

pzelasko requested a review from a team as a code owner September 1, 2026 20:56

github-actions Bot added the ASR label Sep 1, 2026

pzelasko force-pushed the codex/pr16158-rewrite/02-indexed-resumable-data branch from b611e76 to 1fa9921 Compare September 1, 2026 22:29

pzelasko force-pushed the codex/pr16158-rewrite/03-packed-speech-training branch from d707a1b to 42a1f6e Compare September 1, 2026 22:48

copy-pr-bot Bot temporarily deployed to public September 1, 2026 22:50 Inactive

copy-pr-bot Bot had a problem deploying to test September 1, 2026 22:51 Error

github-advanced-security Bot found potential problems Sep 1, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public September 1, 2026 22:53 Inactive

copy-pr-bot Bot temporarily deployed to public September 1, 2026 22:54 Inactive

pzelasko force-pushed the codex/pr16158-rewrite/03-packed-speech-training branch from 42a1f6e to a0eb08d Compare September 1, 2026 23:11

copy-pr-bot Bot temporarily deployed to public September 1, 2026 23:13 Inactive

copy-pr-bot Bot temporarily deployed to test September 1, 2026 23:14 Inactive

copy-pr-bot Bot temporarily deployed to public September 1, 2026 23:16 Inactive

copy-pr-bot Bot temporarily deployed to public September 1, 2026 23:17 Inactive

pzelasko force-pushed the codex/pr16158-rewrite/02-indexed-resumable-data branch from f2946ad to 1979410 Compare September 2, 2026 12:56

pzelasko force-pushed the codex/pr16158-rewrite/03-packed-speech-training branch from a0eb08d to 7eb5294 Compare September 2, 2026 12:56

copy-pr-bot Bot temporarily deployed to public September 2, 2026 12:58 Inactive

copy-pr-bot Bot temporarily deployed to test September 2, 2026 12:59 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 13:02 Inactive

nithinraok reviewed Sep 10, 2026

View reviewed changes

nithinraok previously approved these changes Sep 11, 2026

View reviewed changes

chtruong814 requested changes Sep 11, 2026

View reviewed changes

pzelasko mentioned this pull request Sep 11, 2026

build(speechlm2): add cut-cross-entropy to compiled extras #16241

Open

chtruong814 previously approved these changes Sep 11, 2026

View reviewed changes

nithinraok previously approved these changes Sep 11, 2026

View reviewed changes

pzelasko and others added 16 commits September 11, 2026 21:21

style(speechlm2): format packed training changes

f7c69c2

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

fix(asr): stabilize packed feature normalization

7409c06

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

refactor(speechlm2): break encoder loading import cycles

744d15d

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

fix(speechlm2): stabilize packed CI coverage

012517d

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

refactor(speechlm2): unify packed PEE chunking semantics

98f5e68

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

test(asr): clone experimental PEE via state dict

83cb882

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

test(speechlm2): assert all-empty FSDP parameter grads

1ff09ad

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

style(asr): format packed transformer cleanup

a3c799d

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

fix(speechlm2): retain checkpoint metadata during restore

44c7054

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

refactor(asr): remove obsolete packed encoder paths

bd9c341

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

refactor(asr): isolate packed attention utilities

23ab8f4

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

docs(asr): document packed attention helpers

5adc00c

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

test(speechlm2): stabilize packed perception coverage

919c31b

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

		return flash_attn_varlen_func


		def _forward_sequence_packed_layer(layer, x, **kwargs):

Uh oh!

Conversation

pzelasko commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

feat(speechlm2): add padding-free packed speech training

Summary

Reviewer guide

Attribution

Validation

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

pzelasko commented Sep 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pzelasko commented Sep 1, 2026

Uh oh!

pzelasko commented Sep 2, 2026

Uh oh!

pzelasko commented Sep 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pzelasko commented Sep 11, 2026

Uh oh!

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pzelasko commented Sep 11, 2026

Uh oh!

pzelasko commented Sep 11, 2026

Uh oh!

chtruong814 left a comment

Choose a reason for hiding this comment

Uh oh!

chtruong814 commented Sep 11, 2026

Uh oh!

pzelasko commented Sep 11, 2026

Uh oh!

github-actions Bot commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pzelasko commented Sep 1, 2026 •

edited

Loading